import argparse
import onnx
from google.protobuf.json_format import MessageToDict, MessageToJson
import json

parser = argparse.ArgumentParser(description="Get ONNX model input/output info")
parser.add_argument("-m", "--model", required=True, nargs=1, type=str, help="ONNX file")
args = parser.parse_args()
source_file = args.model[0]

model = onnx.load(source_file)
output_names = [node.name for node in model.graph.output]

input_all = [node.name for node in model.graph.input]
input_initializer = [node.name for node in model.graph.initializer]
net_feed_input = list(set(input_all) - set(input_initializer))

print("================== input/output names ==================")
print("Input names: ", net_feed_input)
print("Output names: ", output_names)
print()

print("================== input info ==========================")
for _input in model.graph.input:
    input_dict = MessageToDict(_input)
    input_name = input_dict["name"]
    if (input_name in net_feed_input):
        print(json.dumps(MessageToDict(_input), sort_keys=False, indent=4))
print()
print("================== output info =========================")
for _output in model.graph.output:
    print(json.dumps(MessageToDict(_output), sort_keys=False, indent=4))
print()

print("================== data type reference ==================")
print(onnx.TensorProto.DataType.items())
print()